package org.finesys.ai.core;

import static org.finesys.ai.core.HeaderNames.ORGANIZATION;
import static org.finesys.ai.core.HeaderNames.USER_AGENT;

import java.io.IOException;
import java.util.HashMap;
import java.util.Map;
import java.util.Objects;

import org.finesys.ai.core.balance.BalanceResponse;
import org.finesys.ai.core.chat.ChatCompletionRequest;
import org.finesys.ai.core.chat.ChatCompletionResponse;
import org.finesys.ai.core.chat.ChatRequest;
import org.finesys.ai.core.completion.CompletionRequest;
import org.finesys.ai.core.completion.CompletionResponse;
import org.finesys.ai.core.models.ModelsResponse;
import org.finesys.ai.core.moderation.ModerationRequest;
import org.finesys.ai.core.moderation.ModerationResponse;
import org.finesys.ai.core.moderation.ModerationResult;
import org.finesys.ai.core.search.SearchApi;
import org.finesys.ai.core.search.SearchRequest;
import org.finesys.ai.core.search.SearchResponse;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import okhttp3.Cache;
import okhttp3.OkHttpClient;
import reactor.core.publisher.Flux;
import retrofit2.Retrofit;
import retrofit2.converter.jackson.JacksonConverterFactory;

/**
 * deepseek模型客户端
 */
public class DeepSeekClient extends OpenAiClient {

    private static final Logger log = LoggerFactory.getLogger(DeepSeekClient.class);
    private final String baseUrl;

    private final String apiVersion;

    private final String model;

    private final OkHttpClient okHttpClient;

    private final OpenAiApi openAiApi;

    private final SearchApi searchApi;

    private final boolean logStreamingResponses;

    private final String systemMessage;

    public DeepSeekClient(String apiKey) {
        this(new Builder().openAiApiKey(apiKey));
    }

    /**
     * DeepSeekClient的构造函数
     *
     * @param serviceBuilder 一个Builder对象，包含构建DeepSeekClient所需的所有配置信息
     * @throws IllegalArgumentException 如果openAiApiKey和azureApiKey同时未定义或同时定义，则抛出此异常
     */
    public DeepSeekClient(Builder serviceBuilder) {
        this.baseUrl = serviceBuilder.baseUrl;
        this.apiVersion = serviceBuilder.apiVersion;
        this.model = serviceBuilder.model;
        this.systemMessage = serviceBuilder.systemMessage;

        OkHttpClient.Builder okHttpClientBuilder = new OkHttpClient.Builder().callTimeout(serviceBuilder.callTimeout)
                .connectTimeout(serviceBuilder.connectTimeout).readTimeout(serviceBuilder.readTimeout)
                .writeTimeout(serviceBuilder.writeTimeout);

        if (serviceBuilder.dispatcher != null) {
            okHttpClientBuilder.dispatcher(serviceBuilder.dispatcher);
        }

        if (serviceBuilder.openAiApiKey == null && serviceBuilder.azureApiKey == null) {
            throw new IllegalArgumentException("openAiApiKey OR azureApiKey must be defined");
        }
        if (serviceBuilder.openAiApiKey != null && serviceBuilder.azureApiKey != null) {
            throw new IllegalArgumentException("openAiApiKey AND azureApiKey cannot both be defined at the same time");
        }
        if (serviceBuilder.openAiApiKey != null) {
            okHttpClientBuilder.addInterceptor(new AuthorizationHeaderInjector(serviceBuilder.openAiApiKey));
        } else {
            okHttpClientBuilder.addInterceptor(new ApiKeyHeaderInjector(serviceBuilder.azureApiKey));
        }

        Map<String, String> headers = new HashMap<>();
        if (serviceBuilder.organizationId != null) {
            headers.put(ORGANIZATION, serviceBuilder.organizationId);
        }
        if (serviceBuilder.userAgent != null) {
            headers.put(USER_AGENT, serviceBuilder.userAgent);
        }
        if (serviceBuilder.customHeaders != null) {
            headers.putAll(serviceBuilder.customHeaders);
        }
        if (!headers.isEmpty()) {
            okHttpClientBuilder.addInterceptor(new GenericHeaderInjector(headers));
        }

        if (serviceBuilder.proxy != null) {
            okHttpClientBuilder.proxy(serviceBuilder.proxy);
        }

        if (serviceBuilder.logRequests) {
            okHttpClientBuilder.addInterceptor(new RequestLoggingInterceptor(serviceBuilder.logLevel));
        }

        if (serviceBuilder.logResponses) {
            okHttpClientBuilder.addInterceptor(new ResponseLoggingInterceptor(serviceBuilder.logLevel));
        }
        this.logStreamingResponses = serviceBuilder.logStreamingResponses;

        this.okHttpClient = okHttpClientBuilder.build();

        Retrofit.Builder retrofitBuilder = new Retrofit.Builder().baseUrl(serviceBuilder.baseUrl).client(okHttpClient);

        if (serviceBuilder.persistTo != null) {
            retrofitBuilder.addConverterFactory(new PersistorConverterFactory(serviceBuilder.persistTo));
        }
        retrofitBuilder.addConverterFactory(JacksonConverterFactory.create(Json.OBJECT_MAPPER));

        this.openAiApi = retrofitBuilder.build().create(OpenAiApi.class);

        // Search API
        OkHttpClient searchOkHttpClient = new OkHttpClient.Builder().callTimeout(serviceBuilder.callTimeout)
                .connectTimeout(serviceBuilder.connectTimeout).readTimeout(serviceBuilder.readTimeout)
                .writeTimeout(serviceBuilder.writeTimeout)
                .addInterceptor(new RequestLoggingInterceptor(serviceBuilder.logLevel))
                .addInterceptor(new ResponseLoggingInterceptor(serviceBuilder.logLevel))
                .addInterceptor(new AuthorizationHeaderInjector(serviceBuilder.searchApiKey)).build();
        Retrofit.Builder searchRetrofitBuilder = new Retrofit.Builder().baseUrl(serviceBuilder.searchEndpoint)
                .client(searchOkHttpClient);
        searchRetrofitBuilder.addConverterFactory(JacksonConverterFactory.create(Json.OBJECT_MAPPER));
        this.searchApi = searchRetrofitBuilder.build().create(SearchApi.class);
    }

    /**
     * 关闭OkHttpClient的相关资源
     * <p>
     * 关闭OkHttpClient的调度器、连接池和缓存资源。
     * 如果缓存资源不为空，则尝试关闭缓存，并捕获并处理可能出现的IOException异常。
     */
    public void shutdown() {
        okHttpClient.dispatcher().executorService().shutdown();

        okHttpClient.connectionPool().evictAll();

        Cache cache = okHttpClient.cache();
        if (cache != null) {
            try {
                cache.close();
            } catch (IOException e) {
                log.error("Failed to close cache", e);
            }
        }
    }

    public static Builder builder() {
        return new Builder();
    }

    public static class Builder extends OpenAiClient.Builder<DeepSeekClient, Builder> {

        private String searchEndpoint = "https://api.deepseek.com/v1/";

        private String searchApiKey;

        private SearchApi searchApi;

        public Builder searchApi(SearchApi searchApi) {
            this.searchApi = searchApi;
            return this;
        }

        public Builder searchApiKey(String searchApiKey) {
            this.searchApiKey = searchApiKey;
            return this;
        }

        public Builder searchEndpoint(String searchEndpoint) {
            this.searchEndpoint = searchEndpoint;
            return this;
        }

        public DeepSeekClient build() {
            return new DeepSeekClient(this);
        }

    }


    @Override
    public BalanceResponse balance() {
        return new RequestExecutor<>(this.openAiApi.balance(new HashMap<>()), r -> r, okHttpClient, null,
                null, BalanceResponse.class, null, logStreamingResponses).execute();
    }


    /**
     * 处理完成请求，支持同步、异步或流式传输。
     *
     * @param context OpenAI客户端上下文信息
     * @param request 完成请求
     * @return SyncOrAsyncOrStreaming对象，包含完成响应
     */
    @Override
    public SyncOrAsyncOrStreaming<CompletionResponse> completion(OpenAiClientContext context,
                                                                 CompletionRequest request) {
        CompletionRequest syncRequest = CompletionRequest.builder().from(request).stream(false).build();

        return new RequestExecutor<>(openAiApi.completions(context.headers(), syncRequest, apiVersion), r -> r,
                okHttpClient, formatUrl("completions"),
                () -> CompletionRequest.builder().from(request).stream(true).build(), CompletionResponse.class, r -> r,
                logStreamingResponses);
    }

    /**
     * 根据上下文和提示完成文本。
     *
     * @param context OpenAI客户端上下文
     * @param prompt  文本提示
     * @return 返回一个包含完成文本的对象，可以是同步、异步或流式传输
     */
    @Override
    public SyncOrAsyncOrStreaming<String> completion(OpenAiClientContext context, String prompt) {
        CompletionRequest request = CompletionRequest.builder().prompt(prompt).build();

        CompletionRequest syncRequest = CompletionRequest.builder().from(request).stream(false).build();

        return new RequestExecutor<>(openAiApi.completions(context.headers(), syncRequest, apiVersion),
                CompletionResponse::text, okHttpClient, formatUrl("completions"),
                () -> CompletionRequest.builder().from(request).stream(true).build(), CompletionResponse.class,
                CompletionResponse::text, logStreamingResponses);
    }

    /**
     * 根据给定的上下文和请求，生成对话补全结果。
     *
     * @param context OpenAiClientContext类型的上下文对象，包含API请求所需的认证信息等。
     * @param request ChatCompletionRequest类型的请求对象，包含模型名称、提示信息等参数。
     * @return SyncOrAsyncOrStreaming<ChatCompletionResponse>类型的结果，表示生成的对话补全结果。
     * @throws NullPointerException 如果请求对象中的模型名称为null且未设置默认模型，则抛出此异常。
     */
    @Override
    public SyncOrAsyncOrStreaming<ChatCompletionResponse> chatCompletion(OpenAiClientContext context,
                                                                         ChatCompletionRequest request) {

        if (Objects.isNull(request.getModel())) {
            request.setModel(this.model);
        }

        ChatCompletionRequest syncRequest = ChatCompletionRequest.builder().from(request).stream(false).build();

        return new RequestExecutor<>(openAiApi.chatCompletions(context.headers(), syncRequest, apiVersion), r -> r,
                okHttpClient, formatUrl("chat/completions"),
                () -> ChatCompletionRequest.builder().from(request).stream(true).build(), ChatCompletionResponse.class,
                r -> r, logStreamingResponses);
    }

    /**
     * 根据用户消息生成对话补全结果
     *
     * @param context     OpenAiClient的上下文信息
     * @param userMessage 用户发送的消息
     * @return 包含对话补全结果的SyncOrAsyncOrStreaming对象
     */
    @Override
    public SyncOrAsyncOrStreaming<String> chatCompletion(OpenAiClientContext context, String userMessage) {
        ChatCompletionRequest request = ChatCompletionRequest.builder().addUserMessage(userMessage).build();

        ChatCompletionRequest syncRequest = ChatCompletionRequest.builder().from(request).stream(false).build();

        return new RequestExecutor<>(openAiApi.chatCompletions(context.headers(), syncRequest, apiVersion),
                ChatCompletionResponse::content, okHttpClient, formatUrl("chat/completions"),
                () -> ChatCompletionRequest.builder().from(request).stream(true).build(), ChatCompletionResponse.class,
                r -> r.choices().get(0).delta().content(), logStreamingResponses);
    }


    /**
     * 与用户进行聊天并完成Flux响应。
     *
     * @param request 请求信息，包含聊天内容和模型信息
     * @return Flux<ChatCompletionResponse> 返回聊天完成响应的Flux流
     */
    @Override
    public Flux<ChatCompletionResponse> chatFluxCompletion(ChatCompletionRequest request) {
        if (Objects.isNull(request.getModel())) {
            request.setModel(this.model);
        }

        return Flux.create(emitter -> {
            this.chatCompletion(new OpenAiClientContext(), request).onPartialResponse(emitter::next)
                    .onComplete(emitter::complete).onError(emitter::error).execute();
        });
    }


    /**
     * 根据用户输入的消息搜索并生成聊天补全响应流。
     *
     * @param userMessage 用户输入的消息
     * @return 聊天补全响应流
     */
    public Flux<ChatCompletionResponse> chatSearchCompletion(String userMessage) {
        SearchResponse searchResponse = new SyncRequestExecutor<>(
                searchApi.webSearch(SearchRequest.builder().enable(true).query(userMessage).build()),
                searchResponse1 -> searchResponse1).execute();
        if (200 == searchResponse.getCode()) {
            String formatted = Utils.format(userMessage, searchResponse.getData().getWebPages().getValues());
            if (Objects.nonNull(formatted)) {
                return this.chatFluxCompletion(formatted);
            }
            return this.chatFluxCompletion(userMessage);
        }

        return this.chatFluxCompletion(userMessage);
    }


    /**
     * 使用用户消息进行流式搜索并完成聊天。
     *
     * @param userMessage 用户输入的消息
     * @return 返回一个SyncOrAsyncOrStreaming<ChatCompletionResponse>对象，包含聊天完成的响应
     */
    public SyncOrAsyncOrStreaming<ChatCompletionResponse> chatSearchStreamingCompletion(String userMessage) {
        SearchResponse searchResponse = new SyncRequestExecutor<>(
                searchApi.webSearch(SearchRequest.builder().enable(true).query(userMessage).build()),
                searchResponse1 -> searchResponse1).execute();
        if (200 == searchResponse.getCode()) {
            String formatted = Utils.format(userMessage, searchResponse.getData().getWebPages().getValues());
            if (Objects.nonNull(formatted)) {
                return this.chatCompletion(new OpenAiClientContext(), ChatCompletionRequest.builder().stream(true)
                        .model(this.model).addUserMessage(formatted).build());
            }
        }
        return this.chatCompletion(new OpenAiClientContext(),
                ChatCompletionRequest.builder().stream(true).model(this.model).addUserMessage(userMessage).build());
    }


    /**
     * 根据用户输入的消息和搜索请求，执行搜索并返回搜索结果
     *
     * @param userMessage   用户输入的消息
     * @param searchRequest 搜索请求对象
     * @return 搜索完成后的响应流
     */
    public Flux<ChatCompletionResponse> chatSearchCompletion(String userMessage, SearchRequest searchRequest) {

        if (Objects.isNull(searchRequest.getQuery())) {
            searchRequest.setQuery(userMessage);
        }

        SearchResponse searchResponse = new SyncRequestExecutor<>(searchApi.webSearch(searchRequest),
                searchResponse1 -> searchResponse1).execute();
        if (200 == searchResponse.getCode()) {
            String formatted = Utils.format(userMessage, searchResponse.getData().getWebPages().getValues());
            if (Objects.nonNull(formatted)) {
                return this.chatFluxCompletion(formatted);
            }
            return this.chatFluxCompletion(userMessage);
        }

        return this.chatFluxCompletion(userMessage);
    }

    /**
     * 根据用户消息异步生成聊天回复流
     *
     * @param userMessage 用户输入的消息
     * @return 一个包含多个ChatCompletionResponse对象的Flux流
     */
    @Override
    public Flux<ChatCompletionResponse> chatFluxCompletion(String userMessage) {
        ChatCompletionRequest.Builder builder = ChatCompletionRequest.builder();

        if (Objects.nonNull(this.systemMessage)) {
            builder.addSystemMessage(this.systemMessage);
        }
        builder.addUserMessage(userMessage);
        return this.chatFluxCompletion(builder.build());
    }

    /**
     * 通过Flux异步返回聊天补全响应
     *
     * @param request 聊天请求
     * @return 包含聊天补全响应的Flux对象
     */
    @Override
    public Flux<ChatCompletionResponse> chatFluxCompletion(ChatRequest request) {
        ChatCompletionRequest.Builder builder = ChatCompletionRequest.builder();
        if (Objects.nonNull(request.getSystemMessage())) {
            builder.addSystemMessage(request.getSystemMessage());
        }
        if (Objects.nonNull(request.getUserMessage())) {
            builder.addUserMessage(request.getUserMessage());
        }
        if (Objects.nonNull(request.getAssistantMessage())) {
            builder.addAssistantMessage(request.getAssistantMessage());
        }
        if (Objects.nonNull(request.getModel())) {
            builder.model(request.getModel());
        }
        if (Objects.nonNull(request.getTemperature())) {
            builder.temperature(request.getTemperature());
        }
        if (Objects.nonNull(request.getMaxTokens())) {
            builder.maxCompletionTokens(request.getMaxTokens());
        }
        if (Objects.nonNull(request.getStopWords())) {
            builder.stop(request.getStopWords());
        }

        builder.stream(request.getStream());
        return this.chatFluxCompletion(builder.build());
    }

    /**
     * 对内容进行审核。
     *
     * @param context OpenAi客户端上下文，包含请求所需的headers等信息
     * @param request 审核请求，包含需要审核的内容等信息
     * @return 返回审核结果的SyncOrAsync对象，其中包含ModerationResponse对象
     */
    @Override
    public SyncOrAsync<ModerationResponse> moderation(OpenAiClientContext context, ModerationRequest request) {
        return new RequestExecutor<>(openAiApi.moderations(context.headers(), request, apiVersion), r -> r);
    }

    /**
     * 对输入内容进行内容审核
     *
     * @param context 客户端上下文，包含请求头等信息
     * @param input   待审核的内容
     * @return SyncOrAsync<ModerationResult> 审核结果，可能同步或异步返回
     */
    @Override
    public SyncOrAsync<ModerationResult> moderation(OpenAiClientContext context, String input) {
        ModerationRequest request = ModerationRequest.builder().input(input).build();

        return new RequestExecutor<>(openAiApi.moderations(context.headers(), request, apiVersion),
                r -> r.results().get(0));
    }

    /**
     * 获取模型列表。
     *
     * @return 包含模型列表的ModelsResponse对象。
     */
    public ModelsResponse models() {

        return new RequestExecutor<>(this.openAiApi.models(new HashMap<>(), apiVersion), r -> r, okHttpClient, null,
                null, ModelsResponse.class, null, logStreamingResponses).execute();
    }

    /**
     * 格式化URL
     *
     * @param endpoint 目标URL的端点部分
     * @return 格式化后的完整URL
     */
    private String formatUrl(String endpoint) {
        return baseUrl + endpoint + apiVersionQueryParam();
    }

    /**
     * 获取API版本查询参数
     *
     * @return API版本查询参数字符串
     */
    private String apiVersionQueryParam() {
        if (apiVersion == null || apiVersion.trim().isEmpty()) {
            return "";
        }
        return "?api-version=" + apiVersion;
    }

}
